from flask import Blueprint, request,Response
from langchain_core.prompts import ChatPromptTemplate

from pojo.result import Result  # 所有的依赖必须从根目录开始引入
#from langchain_community.chat_models import ChatOpenAI
from langchain_openai import ChatOpenAI
#from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from config.Config import llmurl
from config.Template import template
import  requests
import itertools
import json
import time

"""
测试知识库
"""
knowledgebasetest = Blueprint("knowledgebasetest", __name__, url_prefix="/knowledgebasetest")

#搜索知识库
def search_doc(data):
    headers = {
        'accept': 'application/json',
        'Content-Type': 'application/json'
    }
    search_data = {
            "query": data['query'],
            "knowledge_base_name": data['knowledge_base_name'],
            "top_k": data['top_k'],
            "score_threshold": data['score_threshold'],
            "file_name": "",
            "metadata": {}
        }
    resp = requests.post('http://192.168.0.194:7861/knowledge_base/search_docs', json = search_data)
    print("这是搜索知识库返回的数据")
    print(resp.status_code)
    return resp.json()


#生成llmchain
def return_llmchain():
    model = "openai-api"
    temperature = 0.7
    api_key = "EMPTY"
    base_url = "http://192.168.0.194:20000/v1"
    llm = ChatOpenAI(model=model, temperature=temperature, api_key=api_key, base_url=base_url, streaming=True,
                     max_tokens=2048)

    prompt = ChatPromptTemplate.from_template(
        template = template
    )
    llm_chain = prompt | llm
    return llm_chain


#定义两个接口，一个流式输出问答，一个返回查找文档的标题


#流失输出问答
@knowledgebasetest.route('/chat',methods = ['POST'])
def chat():
    data = request.json
    #搜索知识库返回文档数据，以列表的形式返回
    searchlist = search_doc(data)
    query = data['query']
    context = "\n".join([searchdoc["page_content"] for searchdoc in searchlist])
    print(context)
    llm_chain = return_llmchain()
    def return_llm_stream():
        ret = llm_chain.stream({'question': query, 'context': context})
        #for content in itertools.islice(ret, None, None, -1):
        for content in ret:
            #print(content.find("="))
            content = str(content)[9:-1]
            print(content)
            yield content
    return Response(return_llm_stream(), mimetype='text/plain')

#直接返回命中的文档
@knowledgebasetest.route('/returndoc',methods = ['POST'])
def returndoc():

    return vars(Result('200','返回成功','returndoc'))